package org.stragma.apply.modules

import scala.collection.mutable.HashMap

import org.kiama.rewriting._
import org.kiama.rewriting.Rewriter.{all,attempt,bottomup,fail,id,oncetd,rulefs,Term,Strategy,where}

import org.stragma.apply.signatures.Apply._
import org.stragma.apply.signatures._
import org.stragma.apply.util.LogicalOperators._

class Assertions {
	/*
	 * Dynamic rule which might be re-defined at some point
	 */
	var Assertion1 = fail
	val assertion1Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion2 = fail
	val assertion2Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion3 = fail
	val assertion3Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion4 = fail
	val assertion4Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion5 = fail
	val assertion5Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion6 = fail
	val assertion6Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var Assertion7 = fail
	val assertion7Map : HashMap[TExp,Strategy] = new HashMap[TExp,Strategy]()
	
	var UsedInAssertion = fail
	val usedInAssertionMap : HashMap[TVar,Strategy] = new HashMap[TVar,Strategy]()

	/*
	 * Interface: try to employ assertions at this point in the tree
  	 * try-assertion = try(employ-assertion)
	 */
	def try_assertion = attempt(employ_assertion)

	/*
	 * Interface: try to employ assertions at and below this point in the tree
  	 * once-assertion = oncetd(employ-assertion)
	 */
	def once_assertion = oncetd(employ_assertion)

	/*
	 * Interface: try to define an assertion from a conditional expression at the
  	 * current point in the tree   e.g. Leq(Var("x"),Int("1"))
  	 * Expression a is the raw form of the assertion expression
	 */
	def define_assertion = 
		attempt(define_assertion_gt 
			<+ define_assertion_geq 
			<+ define_assertion_equ 
			<+ define_assertion_leq
			<+ define_assertion_neq 
			<+ define_assertion_false
			<+ define_assertion_true)

	/*
	 * Interface: undefine assertions related to a variable which is
	 * at the current point in the tree   e.g. Var("x")
	 */
	def undefine_assertions = where(
		rulefs {
			case v:Var => 
				// Remove the corresponsing assertion
				removeUsedInAssertion(v)
		}
	)

	/*
	 * Kiama does not support the '-' operator for composition of
	 * dynamic rules, so we have supported this using a HashMap.
	 * This particular method supports the - operator
	 */
	private def removeUsedInAssertion(v:Var) = {
		if (usedInAssertionMap.isDefinedAt(v)) {
			if ((usedInAssertionMap(v))(v).isDefined) {
				assertion1Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion1 = fail
				assertion1Map.foreach { entry => Assertion1 = Assertion1 <+ entry }
				
				assertion2Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion2 = fail
				assertion2Map.foreach { entry => Assertion2 = Assertion2 <+ entry }
				
				assertion3Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion3 = fail
				assertion3Map.foreach { entry => Assertion3 = Assertion3 <+ entry }
				
				assertion4Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion4 = fail
				assertion4Map.foreach { entry => Assertion4 = Assertion4 <+ entry }
				
				assertion5Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion5 = fail
				assertion5Map.foreach { entry => Assertion5 = Assertion5 <+ entry }
				
				assertion6Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion6 = fail
				assertion6Map.foreach { entry => Assertion6 = Assertion6 <+ entry }
				
				assertion7Map -= (usedInAssertionMap(v))(v).get.asInstanceOf[TExp]
				Assertion7 = fail
				assertion7Map.foreach { entry => Assertion7 = Assertion7 <+ entry }
			}
		
			usedInAssertionMap -= v
			UsedInAssertion = fail
			usedInAssertionMap.foreach { entry => UsedInAssertion = UsedInAssertion <+ entry }
		}
	}

	/*
	 * Extract all the vars used in an expression, this method
	 * is used for Common Subexpression Eliminate (CSE)
	 */
	def get_vars(t : Term) = {
		var vars = List[TVar]()
	
		def collect = {
			rulefs {
				case v : Var =>
					vars = vars ::: List(v)
					id
				case _ => id
			}
		}
		
		bottomup(collect)(t)
		vars
	}

	/*
	 * See p34 of Program Transformation with Scoped Dynamic Rewrite Rules
  	 * record-assertions(|e) = get-vars ; map({y : ?Var(y); rules(UsedInAssertion :+ Var(y) -> e)})
	 */
	def record_assertions(e : TExp) = rulefs {
		case t =>
			get_vars(t).foreach {
				// This code is to demonstrate how in Kiama we can support the equivalent
				// of the following rule in Stratego
				// Assertion6 :+ Leq... using the deterministic combinator <+
				case (y) =>
					val newStrat = rulefs { case v : Var if (v==y) => e }
					usedInAssertionMap += y -> newStrat					
					UsedInAssertion = UsedInAssertion <+ newStrat
			}
	}
	
	/*
	 * Port of strategy
	 * employ-assertion = Assertion1 <+ Assertion2 <+ Assertion3 <+ Assertion4 <+ Assertion5 <+ Assertion6 <+ Assertion7
	 */
	def employ_assertion = Assertion1 <+ Assertion2 <+ Assertion3 <+ Assertion4 <+ Assertion5 <+ Assertion6 <+ Assertion7

	/*
	 * If we know e > 5 then it follows that
  	 * true: e > 4  e > 5   e >= 4  e >= 5    e != 4  e != 5  
  	 *		so true: e > j for j <= i; e >= j for j <= i; e != j for j <= i
  	 * false: e < 4  e < 5  e <= 4  e <= 5   e = 4  e = 5
  	 *		so false: e < j for j <= i; e <= j for j <= i; e = j for j <= i
	 * define-assertion-gt =
     *		Gt(?e,Int(?i))
     *		; where(<debug-assertions> ["assert gt: ", <id>])
     *		; rules( Assertion3 :+ Gt(e,Int(j)) -> True where (<leqS>(j,i))   )
     *		; rules( Assertion3 :+ Geq(e,Int(j)) -> True where (<leqS>(j,i))  )
     *		; rules( Assertion3 :+ Neq(e,Int(j)) -> True where (<leqS>(j,i))  )
     *		; rules( Assertion3 :+ Lt(e,Int(j)) -> False where (<leqS>(j,i))  )
     *		; rules( Assertion3 :+ Leq(e,Int(j)) -> False where (<leqS>(j,i)) )
     *		; rules( Assertion3 :+ Equ(e,Int(j)) -> False where (<leqS>(j,i)) )
     *		; where(<record-assertions(|e)>e)
	 */
	def define_assertion_gt = rulefs {
		case Gt(e,i:Int) =>
			val newStrat = rulefs {
				case Gt(e,j:Int)  if (leqS(j,i).isDefined) => True()
				case Geq(e,j:Int) if (leqS(j,i).isDefined) => True()
				case Neq(e,j:Int) if (leqS(j,i).isDefined) => True()
				case Lt(e,j:Int) if (leqS(j,i).isDefined) => False()
				case Leq(e,j:Int) if (leqS(j,i).isDefined) => False()
				case Equ(e,j:Int) if (leqS(j,i).isDefined) => False()
			}
			
			assertion3Map += e -> Assertion3	
			Assertion3 = Assertion3 <+ newStrat
			
			(Assertion3 <* where(record_assertions(e)(e)))(Gt(e,i))
	}

	/*
	 * If we know e >= 5 then it follows that
  	 * true: e > 4      e >= 4  e >= 5   e != 4
  	 *		so true: e > j for j < i; e >= j for j <= i; e != j for j < i
  	 * false: e < 4  e < 5   e <= 4      e = 4
  	 *		so false: e < j for j <= i; e <= j for j < i; e != j for j < i
	 * define-assertion-geq =
     *		?Geq(e,Int(i))
     *		; where(<debug-assertions> ["assert geq: ", <id>])
     *		// If there happens to be a Leq assertion already in place with the same
     *		// value, then we actually know equality
     *		; ( build-equality-assertion-geq
     *	   <+ rules( Assertion4 :+ Gt(e,Int(j)) -> True where (<ltS>(j,i))   )
     *	    ; rules( Assertion4 :+ Geq(e,Int(j)) -> True where (<leqS>(j,i)) )
     *	    ; rules( Assertion4 :+ Neq(e,Int(j)) -> True where (<ltS>(j,i))  )
     *	    ; rules( Assertion4 :+ Lt(e,Int(j)) -> False where (<leqS>(j,i)) )
     *	    ; rules( Assertion4 :+ Leq(e,Int(j)) -> False where (<ltS>(j,i)) )
     *	    ; rules( Assertion4 :+ Equ(e,Int(j)) -> False where (<ltS>(j,i)) )
     *	    ; where(<record-assertions(|e)>e) )
	 */
	def define_assertion_geq = rulefs {
		case Geq(e,i:Int) =>
		
			val func = () => {
				assertion4Map += e -> rulefs {
					case Gt(e,j:Int) if (ltS(j,i).isDefined) 	=> True()
					case Geq(e,j:Int) if (leqS(j,i).isDefined) 	=> True()
					case Neq(e,j:Int) if (ltS(j,i).isDefined) 	=> True()
					case Lt(e,j:Int) if (leqS(j,i).isDefined) 	=> False()
					case Leq(e,j:Int) if (ltS(j,i).isDefined) 	=> False()
					case Equ(e,j:Int) if (ltS(j,i).isDefined) 	=> False()
				}
				Assertion4 = Assertion4 <+ assertion4Map(e)
				
				// We need to return the original term passed to the anonymous function
				id
			}
			
			(build_equality_assertion_geq <+ func() <* where(record_assertions(e)(e)))(Geq(e,i))
	}

	/*
	 * build-equality-assertion-geq = ?Geq(e,Int(i)) ; where(!Leq(e,Int(i)) ; employ-assertion ; ?True
     * ; !Equ(e,Int(i)) ; define-assertion-equ )
	 */
	def build_equality_assertion_geq = rulefs {
		case Geq(e,i:Int) =>
			employ_assertion(Leq(e,i)) match {
				case Some(True()) => where(define_assertion_equ(Equ(e,i)))
				case _ => fail
			}
	}

	/*
	 * If we know e <= 5 then it follows that
  	 * true: e < 6        e <= 6  e <= 5    e != 6    
  	 *		so true: e < j for j > i; e <= j for j >= i; e != j for j > i
  	 * false: e > 6  e > 5  e >= 6         e = 6  
  	 *		so false: e > j for j >= i; e >= j for j > i; e = j for j > i
	 */
	def define_assertion_leq = rulefs {
		case Leq(e,i:Int) =>
					
			val func = () => { 
				assertion1Map += e -> rulefs {
					case Gt(e,j:Int) if (ltS(j,i).isDefined) => True()
					case Geq(e,j:Int) if (leqS(j,i).isDefined) => True()
					case Neq(e,j:Int) if (ltS(j,i).isDefined) => True()
					case Lt(e,j:Int) if (leqS(j,i).isDefined) => False()
					case Leq(e,j:Int) if (ltS(j,i).isDefined) => False()
					case Equ(e,j:Int) if (ltS(j,i).isDefined) => False()
				}
				Assertion1 = Assertion1 <+ assertion1Map(e)
				
				// We need to return the original term passed to the anonymous function
				id
			}
			
			// Apply the strategy with immediate effect
			(build_equality_assertion_leq <+ func() <* where(record_assertions(e)(e)))(Leq(e,i))
	}

	/*
	 * build-equality-assertion-leq = ?Leq(e,Int(i)) ; where(!Geq(e,Int(i)) ; employ-assertion ; ?True
     *                      ; !Equ(e,Int(i)) ; define-assertion-equ )
	 */
	def build_equality_assertion_leq = rulefs {
		case Leq(e,i:Int) =>
			employ_assertion(Geq(e,i)) match {
				case Some(True()) => where(define_assertion_equ)(Equ(e,i))
				case _ => fail
			}
	}

	/*
	 * If we know e < 5 then it follows that
  	 * 	true: e < 6  e < 5   e <= 6  e <= 5    e != 6  e != 5  
  	 *		so true: e < j for j >= i; e <= j for j >= i; e != j for j >= i
  	 * false: e > 6  e > 5  e >= 6  e >= 5   e = 6  e = 5
  	 *		so false: e > j for j >= i; e >= j for j >= i; e = j for j >= i
	 * define-assertion-lt =
     *		Lt(?e,Int(?i))
     *		; where(<debug-assertions> ["assert lt: ", <id>])
     *		; rules( Assertion2 :+ Lt(e,Int(j)) -> True where (<geqS>(j,i))   )
     *		; rules( Assertion2 :+ Leq(e,Int(j)) -> True where (<geqS>(j,i))  )
     *		; rules( Assertion2 :+ Neq(e,Int(j)) -> True where (<geqS>(j,i))  )
     *		; rules( Assertion2 :+ Gt(e,Int(j)) -> False where (<geqS>(j,i))  )
     *		; rules( Assertion2 :+ Geq(e,Int(j)) -> False where (<geqS>(j,i)) )
     *		; rules( Assertion2 :+ Equ(e,Int(j)) -> False where (<geqS>(j,i)) )
     *		; where(<record-assertions(|e)>e)
	 */
	def define_assertion_lt = rulefs {
		case Lt(e,i:Int) =>
						
			assertion2Map += e -> rulefs {
				case Lt(e,j:Int) if (geqS(j,i).isDefined) => True()
				case Leq(e,j:Int) if (geqS(j,i).isDefined) =>True()
				case Neq(e,j:Int) if (geqS(j,i).isDefined) => True()
				case Gt(e,j:Int) if (geqS(j,i).isDefined) => False()
				case Geq(e,j:Int) if (geqS(j,i).isDefined) => False()
				case Equ(e,j:Int) if (geqS(j,i).isDefined) => False()
			}	
			Assertion2 = Assertion2 <+ assertion2Map(e)
			
			(Assertion2 <* where(record_assertions(e)(e)))(Lt(e,i))
	}

	/*
	 * Note: A simpler form of dynamic rules with an explicit integer on the left-hand side
  	 * conflicts with the form above, so we have to use a matching form of rule.
  	 * define-assertion-neq =
     *		Neq(?e,Int(?i))
     *			; where(<debug-assertions> ["assert neq: ", <id>])
     *			; rules( Assertion5 :+ Neq(e,Int(j)) -> True where <eq>(j,i)  )
     *			; rules( Assertion5 :+ Equ(e,Int(j)) -> False where <eq>(j,i) )
     *			; where(<record-assertions(|e)>e)
	 */
	def define_assertion_neq = rulefs {
		case Neq(e,i:Int) =>
			
			assertion5Map += e -> rulefs {
				case Neq(e,j:Int) if (eqS(j,i).isDefined) => True()
				case Equ(e,j:Int) if (eqS(j,i).isDefined) => False()
			}	
			Assertion5 = Assertion5 <+ assertion5Map(e)
			
			(Assertion5 <* where(record_assertions(e)(e)))(Neq(e,i))
	}
	
	/*
	 * Port of strategy define-assertion-equ
	 * If e = i then  e REL j is true exactly when i REL j is true
  	 *	define-assertion-equ =
     *		Equ(?e,Int(?i))
     *			; where(<debug-assertions> ["assert equ: ", <id>])
     *			; rules( Assertion6 :+ Lt(e,Int(j)) -> <compare(ltS)>(i,j)        )
     *			; rules( Assertion6 :+ Leq(e,Int(j)) -> <compare(leqS)>(i,j)      )
     *			; rules( Assertion6 :+ Neq(e,Int(j)) -> <compare(not(eq))>(i,j)   )
     *			; rules( Assertion6 :+ Gt(e,Int(j)) -> <compare(gtS)>(i,j)        )
     *			; rules( Assertion6 :+ Geq(e,Int(j)) -> <compare(geqS)>(i,j)      )
     *			; rules( Assertion6 :+ Equ(e,Int(j)) -> <compare(eq)>(i,j)        )
     *			; rules( Assertion7 : e -> Int(i) )
     *			; where(<record-assertions(|e)>e)
	 */
	def define_assertion_equ = rulefs {
		case Equ(e,i:Int) =>
					
			// Update the map for future use when undefining
			assertion6Map += e -> rulefs {
				case Lt(e,j:Int) => if(ltS(i,j).isDefined) True() else False()
				case Leq(e,j:Int) => if (leqS(i,j).isDefined) True() else False()
				case Neq(e,j:Int) => if (neq(i,j).isDefined) True() else False()
				case Gt(e,j:Int) => if (gtS(i,j).isDefined) True() else False()
				case Geq(e,j:Int) => if (geqS(i,j).isDefined) True() else False()
				case Equ(e,j:Int) => if (eqS(i,j).isDefined) True() else False()
			}
			Assertion6 = Assertion6 <+ assertion6Map(e)
			
			assertion7Map += e -> rulefs { case e => i }
			Assertion7 = assertion7Map(e)
			
			(Assertion6 <* Assertion7 <* where(record_assertions(e)(e)))(Equ(e,i))
	}
	
	/*
	 * Port of define-assertion-false(|a) = ?False 
	 */
	def define_assertion_false = rulefs {
		case t : False => id
	}
	
	/*
	 * Port of define-assertion-true(|a) = ?True
	 */
	def define_assertion_true = rulefs {
		case t : True => id
	}
	
}